import pandas as pd
import numpy as np
from database.database import etf_data
import xgboost as xgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, mean_absolute_error
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, confusion_matrix, ConfusionMatrixDisplay, classification_report
from matplotlib.font_manager import FontProperties
import torch.nn as nn
import torch

font = FontProperties(fname=r"C:\Windows\Fonts\simsun.ttc", size=14)

def feature_engineer_LSTM(data):
    data["change"] = data["收盘"].diff(-1)
    features = ["收盘", "最高", "最低", "成交量", "成交额", "均价"]
    X = data[features]
    y = data['change']
    # 数据标准化
    from sklearn.preprocessing import StandardScaler
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # 划分训练集和测试集
    return train_test_split(X_scaled, y, test_size=0.2, shuffle=False)

def feature_engineer(data):    # 加载数据
    # 查看数据
    print(data.head(20))
    data = data[["时间", "开盘", "收盘", "最高", "最低", "成交量", "成交额", "均价"]]
    data["change"] = data["收盘"].diff(-1)

    # 计算移动平均线
    data = data[["开盘", "收盘", "最高", "最低", "成交量", "成交额", "均价", "change"]]
    data['price_5'] = data['收盘'].rolling(window=5).mean()  # 5分钟均价
    data['price_15'] = data['收盘'].rolling(window=15).mean()    # 15分钟均价
    data['deal_price_5'] = data['成交额'].rolling(window=5).mean()   # 5分钟成交额
    data['deal_price_15'] = data['成交额'].rolling(window=15).mean()    # 15分钟成交额
    data['amount_5'] = data['成交量'].rolling(window=5).sum()    # 5分钟成交量
    data['amount_15'] = data['成交量'].rolling(window=15).sum()    # 15分钟成交量

    # 计算相对强弱指标（RSI）
    delta = data['收盘'].diff(1)
    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
    rs = gain / loss
    data['rsi_15'] = 100 - (100 / (1 + rs))

    # 去除前几行缺失值
    data = data.dropna()

    # 定义特征和目标变量
    features = ["收盘", "最高", "最低", "成交量", "成交额", "均价", 'price_5', 'price_15', 'deal_price_5', 'deal_price_15', 'amount_5', 'amount_15', 'rsi_15']
    X = data[features]
    y = data['change']

    # 数据标准化
    from sklearn.preprocessing import StandardScaler
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # 划分训练集和测试集
    return train_test_split(X_scaled, y, test_size=0.2, shuffle=False)


class StockRNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers):
        super(StockRNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, 1)

    def forward(self, x):
        out, _ = self.lstm(x)
        print(out.shape)
        out = self.fc(out[:, ])
        return out


if __name__ == '__main__':
    data_loader = etf_data()
    raw_df = data_loader.load_data(folder="1")
    for code, minutes_df in raw_df.groupby("code"):
        print("当前代码：", code, "数据结构", minutes_df.shape)

        model = "R"

        if model == "R":
            X_train, X_test, y_train, y_test = feature_engineer(minutes_df)
            model = xgb.XGBRegressor(
                objective='reg:squarederror',
                max_depth=5,
                learning_rate=0.1,
                n_estimators=100,
                random_state=42
            )
            #训练模型
            model.fit(X_train, y_train, eval_set=[(X_train, y_train), (X_test, y_test)], verbose=50)

            # 进行预测
            y_pred = model.predict(X_test)

            # 评估模型
            mse = mean_squared_error(y_test, y_pred)
            mae = mean_absolute_error(y_test, y_pred)
            print(f'Mean Squared Error: {mse}')
            print(f'Mean Absolute Error: {mae}')

            # 可视化预测结果
            plt.figure(figsize=(12, 6))
            plt.plot(y_test.index, y_test, label='Actual Price')
            plt.plot(y_test.index, y_pred, label='Predicted Price', linestyle='--')
            plt.legend()
            plt.title('EFT Price Prediction')
            plt.xlabel('Time')
            plt.ylabel('Price')
            plt.show()

            # 获取特征重要性
            feature_importances = model.feature_importances_
            features = ["收盘", "最高", "最低", "成交量", "成交额", "均价",'price_5', 'price_15', 'deal_price_5', 'deal_price_15', 'amount_5', 'amount_15', 'rsi_15']

            # 绘制特征重要性图
            plt.figure(figsize=(10, 6))
            plt.barh(feature_names, feature_importances)
            plt.xlabel('Feature Importance')
            plt.title('Feature Importance Analysis: ' + code)
            plt.show()

        elif model == "C":
            X_train, X_test, y_train, y_test = feature_engineer(minutes_df)
            model = xgb.XGBClassifier(
                max_depth=5,
                learning_rate=0.1,
                n_estimators=100,
                random_state=42
            )
            print(y_train)
            y_train = (y_train >= 0).astype(int)
            y_test = (y_test >= 0).astype(int)
            model.fit(X_train, y_train)

            # 进行预测
            y_pred = model.predict(X_test)
            y_pred_proba = model.predict_proba(X_test)[:, 1]  # 获取预测概率

            # 评估模型
            accuracy = accuracy_score(y_test, y_pred)
            print(y_test, y_pred)
            print(f'模型准确率: {accuracy:.2f}')

            # 绘制 ROC 曲线
            from sklearn.metrics import accuracy_score, roc_curve, auc
            fpr, tpr, thresholds = roc_curve(y_test, y_pred_proba)
            roc_auc = auc(fpr, tpr)

            # 可视化预测结果
            plt.figure(figsize=(12, 6))
            plt.plot(y_test.index, y_test, label='Actual Price')
            plt.plot(y_test.index, y_pred, label='Predicted Price', linestyle='--')
            plt.legend()
            plt.title('EFT Price Prediction ' + code)
            plt.xlabel('Time')
            plt.ylabel('Price')
            plt.show()

            plt.figure()
            plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
            plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
            plt.xlim([0.0, 1.0])
            plt.ylim([0.0, 1.05])
            plt.xlabel('False Positive Rate')
            plt.ylabel('True Positive Rate')
            plt.title('Receiver Operating Characteristic ' + code)
            plt.legend(loc="lower right")
            plt.show()
            feature_names = ["收盘", "最高", "最低", "成交量", "成交额", "均价", 'price_5', 'price_15', 'deal_price_5',
                        'deal_price_15', 'amount_5', 'amount_15', 'rsi_15']
            feature_importances = model.feature_importances_

            # 绘制特征重要性图
            plt.figure(figsize=(10, 6))
            plt.barh(feature_names, feature_importances, fontproperties=font)
            plt.xlabel('Feature Importance')
            plt.title('Feature Importance Analysis: ' + code)
            plt.show()

        elif model == "LSTM":
            X_train, X_test, y_train, y_test = feature_engineer_LSTM(minutes_df)
            X_train_tensor = torch.tensor(X_train, dtype=torch.float32)

            print(X_train)
            import torch
            model = StockRNN(input_size=6, hidden_size=50, num_layers=2)
            criterion = nn.MSELoss()
            optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
            # X_train = X_train[["收盘", "最高", "最低", "成交量", "成交额", "均价"]]
            for epoch in range(50):
                model.train()
                optimizer.zero_grad()
                output = model(X_train_tensor)
                loss = criterion(output, y_train)
                loss.backward()
                optimizer.step()
                print(f'Epoch {epoch + 1}, Loss: {loss.item()}')